package com.yanghui.agent.javassist;

import javassist.*;
import javassist.expr.ExprEditor;
import javassist.expr.MethodCall;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.security.ProtectionDomain;
import java.util.HashSet;
import java.util.Set;

public class JavassistBoot {

    public static class MyMonitorTransformer implements ClassFileTransformer{
        private static final Set<String> classNameSet = new HashSet<>();
        static {
            classNameSet.add("com.yanghui.agent.service.AgentServiceMain");
        }
        @Override
        public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
                                ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
            try {
                String currentClassName = className.replaceAll("/", ".");
                if (!classNameSet.contains(currentClassName)) {
                    return null;
                }
                System.out.println("transform: [" + currentClassName + "]");
                ClassPool pool = ClassPool.getDefault();
                pool.insertClassPath(new LoaderClassPath(loader));
                CtClass ctClass = pool.get(currentClassName);
                CtMethod[] ctMethods = ctClass.getDeclaredMethods();
                for(CtMethod method : ctMethods){
                    if (method.isEmpty()) {
                        continue;
                    }
                    String methodName = method.getName();
                    if ("main".equalsIgnoreCase(methodName)) {
                        continue;
                    }
                    method.addLocalVariable("start", CtClass.longType);
                    method.addLocalVariable("end", CtClass.longType);
                    method.insertBefore("{long start = System.currentTimeMillis();}");
                    method.insertAfter("{long end = System.currentTimeMillis();\n" +
                            "System.out.println(\"" + ctClass.getName() + "." + method.getName() +  " cost \" + (end - start));}");
                }
//                CtBehavior[] methods = ctClass.getDeclaredBehaviors();
//                for (CtBehavior method : methods) {
//                    enhanceMethod(method);
//                }
                return ctClass.toBytecode();
            } catch (Exception e) {
                e.printStackTrace();
            }
            return null;
        }

        private void enhanceMethod(CtBehavior method) throws Exception {
            if (method.isEmpty()) {
                return;
            }
            String methodName = method.getName();
            if ("main".equalsIgnoreCase(methodName)) {
                return;
            }
            final StringBuilder source = new StringBuilder();
            // 前置增强: 打入时间戳
            // 保留原有的代码处理逻辑
            source.append("{")
                    .append("long start = System.currentTimeMillis();\n") //前置增强: 打入时间戳
                    .append("$_ = $proceed($$);\n")              //调用原有代码，类似于method();($$)表示所有的参数
                    .append("System.out.print(\"method:[")
                    .append(methodName).append("]\");")
                    .append("\n")
                    .append("System.out.println(\" cost:[\" +(System.currentTimeMillis() - start)+ \"ms]\");") // 后置增强，计算输出方法执行耗时
                    .append("}");
            ExprEditor editor = new ExprEditor() {
                @Override
                public void edit(MethodCall methodCall) throws CannotCompileException {
                    methodCall.replace(source.toString());
                }
            };
            method.instrument(editor);
        }
    }

    public static void premain(String agentArgs, Instrumentation inst) throws Exception {
        System.out.println("---------------agent javassist start-------------------");

        MyMonitorTransformer monitor = new MyMonitorTransformer();
        inst.addTransformer(monitor);

        System.out.println("---------------agent javassist end-------------------");
    }
}
